import pandas as pd
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif']=['SimHei']
import tushare as ts
print("Tushare当前版本:", ts.__version__)
import os
ts.set_token('268f27f2aa37cf5aed87dd1a21f69cc180049142243cf72031991b88')
pro = ts.pro_api()
#import talib
from datetime import datetime
import time
import csv

# 设置导入数据格式、日期等，股票数据为前复权
def get_data(code,start,end):
    df=pro.daily(ts_code = code,autype='qfq',start_date=start,end_date=end, fields='trade_date,ts_code,open,high,low,close,vol,macd,kdj_k,kdj_d,kdj_j')
    #print("#13 START print(df)")
    #print(df)
    #print("#13 END print(df)")
    df.index = pd.to_datetime(df.trade_date)

    
    #macd, macdsignal, macdhist = talib.MACD(df['close'], fastperiod=12, slowperiod=26, signalperiod=9)
    
    # 错错错错错错错错错错错错错错错错错错错错错错错错错错错错错错
    # # 计算12日EMA
    # ema_12 = df['close'].ewm(span=12, adjust=False).mean()
    # # 计算EMAslow（26日EMA）
    # ema_26 = df['close'].ewm(span=26, adjust=False).mean()
    # df['EMA12']=ema_12
    # df['EMA26']=ema_26
    # #df['MACD']=macd
    # #df['MACD_Signal']=macdsignal
    # #df['MACD_Hist']=macdhist
    # #MACD线 = EMAfast EMAslow
    # df['MACD'] = df['EMA12']-df['EMA26']
    # #信号线（Signal Line）= MACD的9日EMA
    # df['MACD_Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()
    # #柱状图（Histogram）= MACD Signal
    # df['MACD_Histogram'] = df['MACD'] - df['MACD_Signal']
    #print("#21 输出macd:")
    #print(macd)
    #print("#21 输出macd END")
    # 错错错错错错错错错错错错错错错错错错错错错错错错错错错错错错


    #设置把日期作为索引
    #df['ma']=0.0 # Backtrader需要用到
    #df['openinterest']=0.0 #Backtrader需要用到
    #定义两个新的列 ma和openinterest
    #df=df[['open','high','low','close','vol']]
    #重新设置df取值,并返回df
    return df

def check_file_exists_os(file_path):
    return os.path.exists(file_path)

# 下载股票数据,且用CSV保存,保存至指定位置
def acquire_code(code): #只下载一只股票数据,且只用CSV保存  未来可以有自己的数据库
    # csv文件路径    
    path = os.path.join(os.path.join(os.getcwd(),
                                     "csv_file"),code + ".csv")
    # 如果文件存在, 就不处理, 防止重复操作
    if check_file_exists_os(path):
        return 

    #code = '600516.SH'# input("请输入股票代码:\n")
    inp_start = '20170101' #input("请输入开始时间:\n")
    inp_end = datetime.now().strftime("%Y%m%d") #input("请输入结束时间:\n")
    df = get_data(code,inp_start,inp_end)
    #print(df.info())
    #输出统计各列的数据量
    #print("-"*30)
    #分割线
    #print(df.describe())
    #输出常用统计参数
    df.sort_index(inplace=True) #把股票数据按照时间正序排列    
    df.to_csv(path)    
    print(f"已获取{code}数据")
    time.sleep(0.1)# 避免请求频率过高(根据API限制调整)

# 运行函数, 爬取股票数据
# acquire_code('600516.SH')

# 判断是深交所有股票
def is_szse_stock_code(code):
    code_str = str(code).zfill(6)
    # 主板以 00 开头
    # 中小板（已合并到主板）曾以 002 开头
    # 创业板以 30 开头
    # 存托凭证以 184 开头
    # 200开头是B股
    return code_str.startswith(('00', '30', '184','200'))
# 深圳的B股,就是取不到
def is_szse_b_stock_code(code):
    code_str = str(code).zfill(6)
    return code_str.startswith(('200'))
# 北京交易所
def is_bse_stock_code(code):
    code_str = str(code).zfill(6)
    return code_str.startswith(('43', '8','920'))
def is_sse_code(code):
    code_str = str(code).zfill(6)  # 补全为6位
    return code_str.startswith(('60', '688', '900'))

# # 新三版的股票:
# def is_np(code):
#     code_str = str(code).zfill(6)  # 补全为6位
#     return code_str.startswith(('920'))

def read_stock_ids_from_csv():
    stock_ids = []
    try:
        with open('合并.csv', 'r', encoding='utf-8') as file:
            reader = csv.reader(file)
            for row in reader:
                str86 = row[0]
                if is_szse_b_stock_code(row[0]):
                    continue
                if True == is_szse_stock_code(row[0]):
                    str86 = str86+".SZ"                  
                if is_bse_stock_code(row[0]):
                    str86 = str86+".BJ"
                if is_sse_code(row[0]):
                    str86 = str86+".SH"
                # if is_np(row[0]):
                #     str86 = str86+".NP"
                stock_ids.append(str86)
        return stock_ids
    except FileNotFoundError:
        print("文件不存在，请检查文件路径！")


def get_all_stock_to_scv():    
    # 目前没有权限调用这一句
    # # 获取全量股票列表（包含上市状态）
    # stock_list = pro.stock_basic(exchange='', list_status='L', fields='ts_code,symbol,name')
    

    # 全量股票代码列表
    ts_codes = read_stock_ids_from_csv() # stock_list['ts_code'].tolist()
    for idx,ts_code in enumerate(ts_codes):
        if idx>0 and idx%200 ==0:
            time.sleep(60) # 每200次暂停1分钟
        try:
            # 调用你的函数获取单只股票数据
            data = acquire_code(ts_code)
        except Exception as e:
            print(f"获取{ts_code}失败:{e}")

get_all_stock_to_scv()